#include "kernel_operator.h"
#define BUFFER_NUM 2
#define LOOP_BYTE_SIZE 65536

using namespace AscendC;

class ReshapeComm {
    public:
    __aicore__ inline ReshapeComm() {}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y,uint32_t totalLength, TPipe *pipeIn) {

        this->totalLength = totalLength;
        this->LoopSize = LOOP_BYTE_SIZE;
        this->pipe = pipeIn;

        xGm.SetGlobalBuffer((__gm__ int8_t *)x, totalLength);
        yGm.SetGlobalBuffer((__gm__ int8_t *)y, totalLength);

        pipe->InitBuffer(InX, BUFFER_NUM, LOOP_BYTE_SIZE);
        padParams = {false, 0, 0, 0};
        copyParams = {1, 4, 0, 0, 0};
        // printf("Init done! totalLength = %d LoopSize=%d\n",totalLength,LoopSize);
    }
    __aicore__ inline void process() {
        for (uint32_t i = 0; i < totalLength; i+=LoopSize) {
            uint32_t copySize = ((totalLength - i) < LoopSize) ? totalLength - i : LoopSize;
            // printf("i,cnt = %d %d\n",i,copySize);
            copyIn(i, copySize);
            copyOut(i, copySize);
        }

    }
    __aicore__ inline void copyIn(int i, uint32_t cnt) {
        AscendC::LocalTensor<int8_t> inXLocal = InX.AllocTensor<int8_t>();
        // AscendC::DataCopyExtParams copyParams{1,  cnt, 0, 0, 0};
        copyParams.blockLen = cnt;
        // DataCopyPadExtParams<T> padParams = {false, 0, 0, 0};
        AscendC::DataCopyPad(inXLocal, xGm[i], copyParams, padParams);

        InX.EnQue<int8_t>(inXLocal);
    }

    __aicore__ inline void copyOut(int i, uint32_t cnt) {
        AscendC::LocalTensor<int8_t> outXLocal = InX.DeQue<int8_t>();
        // AscendC::DataCopyExtParams copyParams{1,  cnt, 0, 0, 0};
        copyParams.blockLen = cnt;
        AscendC::DataCopyPad(yGm[i], outXLocal, copyParams);
        InX.FreeTensor(outXLocal);
    }

private:
    uint32_t totalLength, totalNum;
    int LoopSize;
    TPipe *pipe;
    // TQue<TPosition::VECIN, BUFFER_NUM> InX;
    TQueBind<TPosition::VECIN, TPosition::VECOUT, BUFFER_NUM> InX;
    DataCopyPadExtParams<int8_t> padParams;
    AscendC::DataCopyExtParams copyParams;
    GlobalTensor<int8_t> yGm, xGm;
};

extern "C" __global__ __aicore__ void reshape(GM_ADDR x, GM_ADDR shape, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    TPipe pipe;

    ReshapeComm op;
    op.Init(x, y, tiling_data.totalLength * (sizeof(DTYPE_X)), &pipe);
    op.process();

}